#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import signal 
import os
import sys
import glob
import cv2
import re
import functools
from PIL import Image
import numpy as np
import argparse
import matplotlib.pyplot as plt
from tqdm import tqdm  # 导入tqdm模块
from pebble import ProcessPool, ProcessExpired
from multiprocessing import cpu_count

def readYuvFile(filename, width, height, stride, channel):
    """
    读取YUV文件并转换为指定通道的图像
    
    Args:
        filename: YUV文件路径
        width: 图像宽度
        height: 图像高度
        stride: 图像步长
        channel: 输出通道类型 ('yuv', 'y', 'u', 'v')
    
    Returns:
        处理后的图像数组
    """
    img_size = int(stride * stride * 1.5)
    yuv = np.zeros(img_size, dtype='uint8', order='C')
    
    with open(filename, "rb") as f:
        f.readinto(yuv)  # 使用readinto提高效率
    
    yuv_img = yuv.reshape((int(stride * 1.5), stride)).astype('uint8')
    
    if channel == 'yuv':
        rgb = cv2.cvtColor(yuv_img, cv2.COLOR_YUV2BGR_NV21)
        return rgb[0:height, 0:width]
    else:
        if channel == 'y':
            channel_img = yuv_img[:height, :width]
        elif channel == 'u':
            u_start = stride * stride
            u_data = yuv[u_start:u_start + stride * stride // 4]
            channel_img = u_data.reshape((height // 2, width // 2))
            channel_img = cv2.resize(channel_img, (width, height), interpolation=cv2.INTER_NEAREST)
        elif channel == 'v':
            v_start = stride * stride + stride * stride // 4
            v_data = yuv[v_start:v_start + stride * stride // 4]
            channel_img = v_data.reshape((height // 2, width // 2))
            channel_img = cv2.resize(channel_img, (width, height), interpolation=cv2.INTER_NEAREST)
        else:
            raise ValueError(f"Unsupported channel: {channel}")
        
        return channel_img
            

def readRgbFile(filename, channel):
    """
    读取RGB图像文件并转换为指定通道
    
    Args:
        filename: 图像文件路径
        channel: 输出通道类型 ('yuv', 'y', 'u', 'v')
    
    Returns:
        处理后的图像数组
    """
    try:
        img = Image.open(filename).convert('RGB')
        bgr_img = np.array(img)
        bgr_img = bgr_img[:,:,::-1]  # 将RGB转换为BGR
        
        if channel == 'yuv':
            return bgr_img
        else:
            yuv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2YUV)
            if channel == 'y':
                channel_img = yuv_img[:,:,0]
            elif channel == 'u':
                channel_img = yuv_img[:,:,1]
            elif channel == 'v':
                channel_img = yuv_img[:,:,2]
            else:
                raise ValueError(f"Unsupported channel: {channel}")
            
            return channel_img
            
    except Exception as e:
        print(f"Error reading RGB file {filename}: {e}")
        return None

def create_output_directory(output_path):
    """
    创建输出目录
    
    Args:
        output_path: 输出路径
    """
    if output_path and not os.path.exists(output_path):
        try:
            os.makedirs(output_path)
            print(f"Created output directory: {output_path}")
        except Exception as e:
            print(f"Error creating output directory {output_path}: {e}")
            return False
    return True

def get_file_list(input_path, suffixes):
    """
    获取文件列表
    
    Args:
        input_path: 输入路径（文件或目录）
        suffixes: 支持的文件后缀列表
    
    Returns:
        文件路径列表
    """
    file_paths = []
    
    if os.path.isfile(input_path):
        file_paths.append(input_path)
    elif os.path.isdir(input_path):
        try:
            all_files = os.listdir(input_path)
            file_paths = [
                os.path.join(input_path, f) 
                for f in all_files 
                if f.lower().endswith(tuple(suffix.lower() for suffix in suffixes))
            ]
            file_paths.sort()  # 按文件名排序
        except Exception as e:
            print(f"Error reading directory {input_path}: {e}")
    else:
        print(f"Input path does not exist: {input_path}")
    
    return file_paths

def save_image(image_array, output_path, overwrite=False):
    """
    保存图像到文件
    
    Args:
        image_array: 图像数组
        output_path: 输出文件路径
        overwrite: 是否覆盖已存在的文件
    
    Returns:
        保存是否成功
    """
    if os.path.exists(output_path) and not overwrite:
        print(f'File {output_path} already exists, skipping (use --overwrite to overwrite)')
        return False
    
    try:
        # 确保图像数组是正确的格式
        if len(image_array.shape) == 3:
            # 彩色图像，转换BGR到RGB
            rgb_image = image_array[:,:,::-1]
        else:
            # 灰度图像
            rgb_image = image_array
        
        # 保存图像
        Image.fromarray(rgb_image).save(output_path)
        print(f'Image saved to {output_path}')
        return True
        
    except Exception as e:
        print(f"Error saving image to {output_path}: {e}")
        return False

def display_image(image_array, window_name, zoom=1.0):
    """
    显示图像
    
    Args:
        image_array: 图像数组
        window_name: 窗口名称
        zoom: 缩放比例
    """
    try:
        if zoom != 1.0:
            new_width = int(image_array.shape[1] * zoom)
            new_height = int(image_array.shape[0] * zoom)
            show_img = cv2.resize(image_array, (new_width, new_height), interpolation=cv2.INTER_AREA)
        else:
            show_img = image_array
        
        cv2.imshow(window_name, show_img)
        
    except Exception as e:
        print(f"Error displaying image {window_name}: {e}")

def yuv_to_png(file_path, output_path, width, height, stride, channel, zoom, overwrite):
    try:
        # 判断文件类型并读取
        ext = os.path.splitext(file_path)[1].lower()
        if ext in ('.png', '.jpg', '.jpeg'):
            image_array = readRgbFile(file_path, channel)
        else:
            image_array = readYuvFile(file_path, width, height, stride, channel)
        
        if image_array is None:
            print(f"Failed to load image from {file_path}, skipping...")
            return None
        # 保存或显示图像
        if output_path is not None:
            # 保存模式
            if os.path.isdir(output_path):
                # 输入是目录，生成输出文件名
                name = os.path.splitext(os.path.basename(file_path))[0]
                output_path = os.path.join(output_path, f'{name}.png')
            save_image(image_array, output_path, overwrite)
        return image_array
    except Exception as e:
        print(f"Error processing file {file_path}: {e}")
        return None

def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='Process YUV/RGB image files and convert them to PNG format or display them.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s input.yuv -o output.png
  %(prog)s input_dir -o output_dir --overwrite
  %(prog)s input.yuv -c y -z 0.5
  %(prog)s input_dir -s yuv raw yuv420 -l 50
        """
    )
    
    # 必需参数
    parser.add_argument("filepath", help="YUV/RGB image file or directory path")
    
    # 图像参数
    parser.add_argument('-iw', '--width', type=int, default=2460, 
                       help='Image width (default: 2460)')
    parser.add_argument('-ih', '--height', type=int, default=2460, 
                       help='Image height (default: 2460)')
    parser.add_argument('-is', '--stride', type=int, default=2496, 
                       help='Image stride (default: 2496)')
    
    # 显示参数
    parser.add_argument('-z', '--zoom', type=float, default=1.0, 
                       help='Image display scale factor (default: 1.0)')
    parser.add_argument('-c', '--channel', type=str, default='yuv', 
                       choices=['yuv','y','u','v'], 
                       help='Image channel to display/save (default: yuv)')
    
    # 输出参数
    parser.add_argument('-o', '--output', type=str, 
                       help='Output file or directory path')
    parser.add_argument('-ow','--overwrite', action='store_true', default=False, 
                       help='Overwrite existing output files (default: False)')
    
    # 处理参数
    parser.add_argument('-n', '--show_max_number', type=int, default=100, 
                       help='Maximum number of images to display (default: 100)')
    
    parser.add_argument('-m','--multi_process', action='store_true', default=False, help='Enable multi-processing mode')
    parser.add_argument("-s", "--suffixes", nargs='+', 
                       default=['yuv', 'raw', 'yuv420', 'png', 'jpg', 'jpeg'], 
                       type=str, help="Supported file suffixes")
    
    args = parser.parse_args()
    
    # 验证输入路径
    if not os.path.exists(args.filepath):
        print(f"Error: Input path does not exist: {args.filepath}")
        return 1
    
    # 创建输出目录
    if args.output:
        output_dir = os.path.dirname(args.output) if os.path.isfile(args.filepath) else args.output
        if output_dir and not create_output_directory(output_dir):
            return 1
    
    # 获取文件列表
    file_paths = get_file_list(args.filepath, args.suffixes)
    
    if not file_paths:
        print("No valid files found to process.")
        return 1
    print(f"Found {len(file_paths)} files to process.")
    if args.multi_process:
        print("Multi-processing mode.")
        with ProcessPool(max_workers=min(len(file_paths), cpu_count())) as pool:
            futures = []
            try:
                # 提交所有转换任务
                for file_path in file_paths:
                    future = pool.schedule(
                        yuv_to_png, 
                        args=(file_path, args.output, args.width, args.height, args.stride, args.channel, args.zoom, args.overwrite),
                        timeout=300  # 单个任务超时时间（秒）
                    )
                    futures.append(future)
                for future in futures:
                    result = future.result()  
            except KeyboardInterrupt:
                print("\n用户中断请求，正在取消所有任务...")
                # 取消所有未完成的任务
                for future in futures:
                    if not future.done():
                        future.cancel()
                # 等待已启动的任务完成
                pool.stop()
                pool.join(timeout=5)
                print("所有任务已终止")
                sys.exit(1)
    else:
        print("Single-threaded mode.")
        processed_count = 0
        display_count = 0
        # 使用进度条（仅当文件数量>=2时）
        iterator = tqdm(file_paths, desc="Processing files") if len(file_paths) >= 2 else file_paths
        for file_path in iterator:
            try:
                image_array = yuv_to_png(file_path, args.output, args.width, args.height, args.stride, args.channel, args.zoom, args.overwrite)
                processed_count += 1
                # 保存或显示图像
                if args.output is None:
                    # 显示模式
                    if display_count < args.show_max_number:
                        display_count += 1
                        window_name = os.path.basename(file_path)
                        display_image(image_array, window_name, args.zoom)
            except Exception as e:
                print(f"Error processing file {file_path}: {e}")
                continue
        
        # 显示处理结果
        if args.output is not None:
            print(f"\nProcessing completed. {processed_count} files processed successfully.")
        else:
            print(f"\nDisplaying {display_count} images. Press any key to close all windows.")
            cv2.waitKey(0)
            cv2.destroyAllWindows()
    
    return 0

if __name__ == '__main__':
    try:
        exit_code = main()
        exit(exit_code)
    except KeyboardInterrupt:
        print("\nOperation cancelled by user.")
        cv2.destroyAllWindows()
        exit(1)
    except Exception as e:
        print(f"Unexpected error: {e}")
        cv2.destroyAllWindows()
        exit(1)
